import numpy as np
import math

def R2q(R):
    """
    R
    q= [x, y, z, w]
    """
        
    w = R[0,0]+R[1,1]+R[2,2]+1
    x = R[0,0]-R[1,1]-R[2,2]+1
    y = -R[0,0]+R[1,1]-R[2,2]+1
    z = -R[0,0]-R[1,1]+R[2,2]+1
    
    q = np.array([w,x,y,z])
    index = np.argmax(q)
    # assert q[index]>0, "max(q) > 0"
    if q[index]<=0:
        return np.array([0,0,0,1],dtype=np.float32)

    q[index] = math.sqrt(q[index]) / 2

    if index==0:
        q0 = q[index]
        q1 = (R[2,1]-R[1,2]) / (4*q0)
        q2 = (R[0,2]-R[2,0]) / (4*q0)
        q3 = (R[1,0]-R[0,1]) / (4*q0)
    elif index==1:
        q1 = q[index]
        q0 = (R[2,1]-R[1,2]) / (4*q1)
        q2 = (R[0,1]+R[1,0]) / (4*q1)
        q3 = (R[2,0]+R[0,2]) / (4*q1)
    elif index==2:
        q2 = q[index]
        q0 = (R[0,2]-R[2,0]) / (4*q2)
        q1 = (R[0,1]+R[1,0]) / (4*q2)
        q3 = (R[1,2]+R[2,1]) / (4*q2)
    elif index==3:
        q3 = q[index]
        q0 = (R[1,0]-R[0,1]) / (4*q3)
        q1 = (R[2,0]+R[0,2]) / (4*q3)
        q2 = (R[1,2]+R[2,1]) / (4*q3)
    else:
        raise ValueError('index error:' +  str(index))
    return np.array([q1,q2,q3,q0],dtype=np.float32)

def q2R(q):
    """
    q: q = [x, y, z, w]
    R
    """
    q = q[[3,0,1,2]]  # (x,y,z,w) -> (w,x,y,z)
    R = np.zeros((3, 3), dtype=float)

    R[0, 0] = 1 - 2 * (q[2] ** 2 + q[3] ** 2)
    R[1, 1] = 1 - 2 * (q[1] ** 2 + q[3] ** 2)
    R[2, 2] = 1 - 2 * (q[1] ** 2 + q[2] ** 2)

    R[0, 1] = 2 * (q[1] * q[2] - q[0] * q[3])
    R[1, 0] = 2 * (q[1] * q[2] + q[0] * q[3])

    R[0, 2] = 2 * (q[1] * q[3] + q[0] * q[2])
    R[2, 0] = 2 * (q[1] * q[3] - q[0] * q[2])

    R[1, 2] = 2 * (q[2] * q[3] - q[0] * q[1])
    R[2, 1] = 2 * (q[2] * q[3] + q[0] * q[1])

    return R